import math
import numpy as np
from scipy.spatial.transform import Rotation

def coord_rotation_matrix(ra, dec):
    '''
    以J2000坐标系为基准，计算赤经赤纬为(ra, dec)恒星坐标系的旋转矩阵
    定义z轴为以合束器为原点，指向观测恒星的单位向量z_unit
    x轴x_unit由z_unit叉乘(0., 0., 1.)得到
    y轴y_unit由z_unit叉乘x_unit得到
    单位向量按列向量排列，形成的3x3矩阵即为旋转矩阵。
    '''
    if dec == 90.:
        print("Are you sure you want to detect the star with dec=90?")
        return np.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]])

    ra = ra / 180. * np.pi
    dec = dec / 180. * np.pi
    ref_z_unit = np.array([0., 0., 1.], dtype=np.float64)
    z_unit = np.array([np.cos(dec) * np.cos(ra), np.cos(dec) * np.sin(ra), np.sin(dec)], dtype=np.float64)
    x_unit = np.cross(z_unit, ref_z_unit)
    y_unit = np.cross(z_unit, x_unit)

    x_unit = x_unit / np.linalg.norm(x_unit)
    y_unit = y_unit / np.linalg.norm(y_unit)

    return np.mat(np.array([x_unit, y_unit, z_unit], dtype=np.float64)).T
    #return np.transpose(np.array([x_unit, y_unit, z_unit], dtype=np.float64))
    
def sate_in_origin_coord(baseline = 30., form_theta = np.pi/8., form_rotate = np.pi / 20):
    if type(form_rotate) == float:
        form_rotate = [form_rotate]
    points_in_coord = np.zeros([4, 3], dtype=np.float64)
    rotate_index = [0., 1., 1., 0.]
    yaxis_index = [1., 1., -1., -1.]
    for i in range(4):
        radian = np.power(-1, i)*form_theta + rotate_index[i]*np.pi + form_rotate[0]
        #points_in_coord[i,0] = baseline / 2. / np.tan(radian)
        #points_in_coord[i,1] = baseline / 2. * yaxis_index[i]
        points_in_coord[i,0] = baseline * np.cos(radian)
        points_in_coord[i,1] = baseline * np.sin(radian)

    return points_in_coord


def sate_in_new_coord(coord_star1=[100., 20.], baseline1=30., coord_star2=[350., -30], form_theta=np.pi/8.):

    rotate_matrix1 = coord_rotation_matrix(coord_star1[0], coord_star1[1])
    rotate_matrix2 = coord_rotation_matrix(coord_star2[0], coord_star2[1])

    origin_sate_points = sate_in_origin_coord(baseline1, form_theta=form_theta)
    print(origin_sate_points)

    rotate1 = Rotation.from_matrix(rotate_matrix1)
    rotate2 = Rotation.from_matrix(rotate_matrix2)
    rotate = rotate2.inv() * rotate1

    sate_after_rotate = rotate.apply(origin_sate_points)
    return sate_after_rotate

def sum_distance(form_rotate:np.float64, sate_coords:np.ndarray, baseline:np.float64, form_theta = np.pi/4.):
    new_sate_coords = sate_in_origin_coord(baseline=baseline, form_theta=form_theta, form_rotate=form_rotate)
    distance_vec = np.linalg.norm(new_sate_coords-sate_coords, axis=1)
    return np.sum(distance_vec)

def trace_planning(coord_star1=[100., 20.], baseline1=30., coord_star2=[350., -30], baseline2=50, form_theta=np.pi/4.):

    sate_after_rotate = sate_in_new_coord(coord_star1=coord_star1, baseline1=baseline1, coord_star2=coord_star2)
    print(sate_after_rotate)
    sum_distance(np.pi/8., sate_after_rotate, baseline2, form_theta=np.pi/16)

    from scipy.optimize import minimize
    form_result = minimize(fun=sum_distance, x0=(0.), args=(sate_after_rotate, baseline2, form_theta), \
                           method = "L-BFGS-B", bounds=((-np.pi, np.pi),))
    print(form_result.x)
    print(form_result.fun)


def main():
    trace_planning()

if __name__ == "__main__":
    main()